package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"socks5_proxy/src/config"
	"socks5_proxy/src/logger"
	"socks5_proxy/src/util"
)

var conf *config.Config

const (
	version = "1.0.1"
)

func printBanner() {
	fmt.Println("       ____    \n  ___ / __/___ \n (_-</__ \\/ _ \\\n/___/____/ .__/\n        /_/    \n")
}

func printVersion() {
	fmt.Println("version:", version)
}

func init() {
	printBanner()
	printVersion()

	err := config.LoadConfig()
	if err != nil {
		fmt.Println("[!] Load configure failed")
		os.Exit(0)
	}

	conf = config.GetConfig()
	logger.InitLogger(conf.LogFile, conf.LogLevel)
	logger.Infof("log system initialized")
}

func main() {
	server, err := net.Listen("tcp", conf.ServiceAddr)
	if err != nil {
		logger.Fatalf("listen failed: %v", err)
		return
	}

	fmt.Println("[+] proxy has been listening...")
	logger.Infof("proxy has been listening on %s", conf.ServiceAddr)

	for {
		client, err := server.Accept()
		if err != nil {
			logger.Fatalf("accept failed: %v", err)
			continue
		}
		go handler(client)
	}
}

func handler(client net.Conn) {
	if err := socks5Auth(client); err != nil {
		logger.Errorf("auth error: %v", err)
		client.Close()
		return
	}

	target, err := socks5Connect(client)
	if err != nil {
		logger.Errorf("connect error: %v", err)
		client.Close()
		return
	}

	socks5Forward(client, target)
}

func socks5Auth(client net.Conn) (err error) {
	buf := make([]byte, 256)

	// read the header.
	n, err := io.ReadFull(client, buf[:2])
	if n != 2 {
		logger.Errorf("bad socks header： %v", err.Error())
		return util.ErrBadHeader
	}

	// read the version and methods.
	ver, nMethods := int(buf[0]), int(buf[1])
	if ver != 5 {
		logger.Errorf("the version is not socks5")
		return util.ErrInvalidVersion
	}

	// read the method list.
	n, err = io.ReadFull(client, buf[:nMethods])
	if n != nMethods {
		return util.ErrBadMethod
	}

	// write back the response that specified without authority.
	n, err = client.Write([]byte{0x05, 0x00})
	if n != 2 || err != nil {
		return util.ErrAuthResponse
	}

	return nil
}

func socks5Connect(client net.Conn) (net.Conn, error) {
	buf := make([]byte, 256)

	n, err := io.ReadFull(client, buf[:4])
	if n != 4 {
		return nil, util.ErrBadHeader
	}

	ver, cmd, _, atyp := buf[0], buf[1], buf[2], buf[3]
	if ver != 5 || cmd != 1 {
		return nil, util.ErrInvalidVersion
	}

	addr := ""
	switch atyp {
	case 1:
		n, err = io.ReadFull(client, buf[:4])
		if n != 4 {
			return nil, util.ErrInvalidIPv4
		}
		addr = fmt.Sprintf("%d.%d.%d.%d", buf[0], buf[1], buf[2], buf[3])

	case 3:
		n, err = io.ReadFull(client, buf[:1])
		if n != 1 {
			return nil, util.ErrInvalidHost
		}
		addrLen := int(buf[0])

		n, err = io.ReadFull(client, buf[:addrLen])
		if n != addrLen {
			return nil, util.ErrInvalidHost
		}
		addr = string(buf[:addrLen])

	case 4:
		return nil, util.ErrIPv6NotSupported

	default:
		return nil, util.ErrInvalidAtyp
	}

	n, err = io.ReadFull(client, buf[:2])
	if n != 2 {
		return nil, util.ErrInvalidPort
	}

	port := binary.BigEndian.Uint16(buf[:2])
	dest, err := net.Dial("tcp", fmt.Sprintf("%s:%d", addr, port))
	if err != nil {
		return nil, util.ErrDialDest
	}

	// write back the response.
	n, err = client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
	if err != nil {
		dest.Close()
		return nil, util.ErrResponse
	}

	return dest, nil
}

func socks5Forward(client, target net.Conn) {
	forward := func(src, dest net.Conn) {
		defer src.Close()
		defer dest.Close()
		io.Copy(src, dest)
	}
	go forward(client, target)
	go forward(target, client)
}
